#!/bin/bash
# ========================================
# RedFire PostgreSQL扩展初始化脚本
# 版本: v5.0.0 (2025-09-16)
# 功能: 安装和配置必要的扩展
# ========================================

set -e

echo "=========================================="
echo "开始安装RedFire PostgreSQL扩展..."
echo "=========================================="

# 等待PostgreSQL完全启动
echo "等待PostgreSQL服务启动..."
while ! pg_isready -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" -h localhost; do
    echo "等待PostgreSQL启动中..."
    sleep 2
done
echo "PostgreSQL服务已启动！"

# 连接到数据库并安装扩展
psql -v ON_ERROR_STOP=1 --username "${POSTGRES_USER}" --dbname "${POSTGRES_DB}" <<-EOSQL

-- ========================================
-- 安装核心扩展
-- ========================================

-- 安装TimescaleDB扩展 (时序数据库)
CREATE EXTENSION IF NOT EXISTS timescaledb CASCADE;
\echo 'TimescaleDB扩展安装完成'

-- 安装UUID扩展 (UUID生成)
CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
\echo 'UUID扩展安装完成'

-- 安装统计扩展 (查询统计)
CREATE EXTENSION IF NOT EXISTS pg_stat_statements;
\echo 'pg_stat_statements扩展安装完成'

-- 安装数组扩展 (数组操作)
CREATE EXTENSION IF NOT EXISTS intarray;
\echo 'intarray扩展安装完成'

-- 安装JSON扩展 (JSON索引)
CREATE EXTENSION IF NOT EXISTS btree_gin;
\echo 'btree_gin扩展安装完成'

-- 安装JSON操作扩展
CREATE EXTENSION IF NOT EXISTS btree_gist;
\echo 'btree_gist扩展安装完成'

-- 安装哈希索引扩展
CREATE EXTENSION IF NOT EXISTS bloom;
\echo 'bloom扩展安装完成'

-- 安装模糊匹配扩展
CREATE EXTENSION IF NOT EXISTS fuzzystrmatch;
\echo 'fuzzystrmatch扩展安装完成'

-- 安装三角函数扩展
CREATE EXTENSION IF NOT EXISTS tablefunc;
\echo 'tablefunc扩展安装完成'

-- 安装加密扩展
CREATE EXTENSION IF NOT EXISTS pgcrypto;
\echo 'pgcrypto扩展安装完成'

-- ========================================
-- 安装可选扩展 (如果可用)
-- ========================================

-- 尝试安装PostGIS扩展 (地理信息系统)
DO \$\$
BEGIN
    CREATE EXTENSION IF NOT EXISTS postgis;
    RAISE NOTICE 'PostGIS扩展安装完成';
EXCEPTION
    WHEN OTHERS THEN
        RAISE NOTICE 'PostGIS扩展不可用，跳过安装: %', SQLERRM;
END
\$\$;

-- 尝试安装全文搜索扩展
DO \$\$
BEGIN
    CREATE EXTENSION IF NOT EXISTS pg_trgm;
    RAISE NOTICE 'pg_trgm扩展安装完成';
EXCEPTION
    WHEN OTHERS THEN
        RAISE NOTICE 'pg_trgm扩展不可用，跳过安装: %', SQLERRM;
END
\$\$;

-- ========================================
-- 配置TimescaleDB
-- ========================================

-- 显示TimescaleDB版本信息
SELECT extname, extversion 
FROM pg_extension 
WHERE extname = 'timescaledb';

-- 显示TimescaleDB许可证信息
SHOW timescaledb.license;

-- ========================================
-- 配置pg_stat_statements
-- ========================================

-- 重置统计信息
SELECT pg_stat_statements_reset();

-- 显示pg_stat_statements配置
SELECT name, setting, unit, short_desc 
FROM pg_settings 
WHERE name LIKE 'pg_stat_statements%'
ORDER BY name;

-- ========================================
-- 创建自定义函数
-- ========================================

-- 创建获取数据库大小的函数
CREATE OR REPLACE FUNCTION get_database_size_mb()
RETURNS TABLE(
    database_name text,
    size_mb numeric
) AS \$\$
BEGIN
    RETURN QUERY
    SELECT 
        datname::text,
        round((pg_database_size(datname) / 1024.0 / 1024.0)::numeric, 2)
    FROM pg_database
    WHERE datname = current_database();
END;
\$\$ LANGUAGE plpgsql;

-- 创建获取表大小的函数
CREATE OR REPLACE FUNCTION get_table_sizes()
RETURNS TABLE(
    schema_name text,
    table_name text,
    total_size_mb numeric,
    table_size_mb numeric,
    index_size_mb numeric,
    toast_size_mb numeric
) AS \$\$
BEGIN
    RETURN QUERY
    SELECT 
        schemaname::text,
        tablename::text,
        round((pg_total_relation_size(schemaname||'.'||tablename) / 1024.0 / 1024.0)::numeric, 2),
        round((pg_relation_size(schemaname||'.'||tablename) / 1024.0 / 1024.0)::numeric, 2),
        round(((pg_total_relation_size(schemaname||'.'||tablename) - pg_relation_size(schemaname||'.'||tablename)) / 1024.0 / 1024.0)::numeric, 2),
        round((COALESCE(pg_total_relation_size(schemaname||'.'||tablename||'_toast'), 0) / 1024.0 / 1024.0)::numeric, 2)
    FROM pg_tables
    WHERE schemaname NOT IN ('information_schema', 'pg_catalog', 'pg_toast', 'timescaledb_information', 'timescaledb_internal')
    ORDER BY pg_total_relation_size(schemaname||'.'||tablename) DESC;
END;
\$\$ LANGUAGE plpgsql;

-- 创建获取连接信息的函数
CREATE OR REPLACE FUNCTION get_connection_stats()
RETURNS TABLE(
    total_connections bigint,
    active_connections bigint,
    idle_connections bigint,
    idle_in_transaction bigint,
    waiting_connections bigint
) AS \$\$
BEGIN
    RETURN QUERY
    SELECT 
        count(*)::bigint,
        count(*) FILTER (WHERE state = 'active')::bigint,
        count(*) FILTER (WHERE state = 'idle')::bigint,
        count(*) FILTER (WHERE state = 'idle in transaction')::bigint,
        count(*) FILTER (WHERE wait_event_type IS NOT NULL)::bigint
    FROM pg_stat_activity;
END;
\$\$ LANGUAGE plpgsql;

-- 创建获取慢查询的函数
CREATE OR REPLACE FUNCTION get_slow_queries(limit_count integer DEFAULT 10)
RETURNS TABLE(
    query_text text,
    calls bigint,
    total_time numeric,
    mean_time numeric,
    rows_returned bigint,
    rows_affected bigint
) AS \$\$
BEGIN
    RETURN QUERY
    SELECT 
        query::text,
        calls::bigint,
        round((total_exec_time)::numeric, 2),
        round((mean_exec_time)::numeric, 2),
        rows::bigint,
        (calls)::bigint
    FROM pg_stat_statements
    WHERE query NOT LIKE '%pg_stat_statements%'
        AND query NOT LIKE '%information_schema%'
    ORDER BY total_exec_time DESC
    LIMIT limit_count;
END;
\$\$ LANGUAGE plpgsql;

-- 创建数据库健康检查函数
CREATE OR REPLACE FUNCTION database_health_check()
RETURNS TABLE(
    check_name text,
    status text,
    value text,
    description text
) AS \$\$
BEGIN
    -- 检查数据库连接
    RETURN QUERY
    SELECT 
        'database_connections'::text,
        CASE WHEN count(*) < (SELECT setting::int FROM pg_settings WHERE name = 'max_connections') * 0.8
             THEN 'OK' ELSE 'WARNING' END::text,
        count(*)::text,
        'Current database connections'::text
    FROM pg_stat_activity;
    
    -- 检查长时间运行的查询
    RETURN QUERY
    SELECT 
        'long_running_queries'::text,
        CASE WHEN count(*) = 0 THEN 'OK' ELSE 'WARNING' END::text,
        count(*)::text,
        'Queries running longer than 5 minutes'::text
    FROM pg_stat_activity
    WHERE state = 'active' 
        AND now() - query_start > interval '5 minutes';
    
    -- 检查死锁
    RETURN QUERY
    SELECT 
        'deadlocks'::text,
        'INFO'::text,
        COALESCE(sum(deadlocks), 0)::text,
        'Total deadlocks detected'::text
    FROM pg_stat_database
    WHERE datname = current_database();
    
    -- 检查磁盘使用
    RETURN QUERY
    SELECT 
        'database_size'::text,
        'INFO'::text,
        pg_size_pretty(pg_database_size(current_database())),
        'Current database size'::text;
    
END;
\$\$ LANGUAGE plpgsql;

-- ========================================
-- 验证扩展安装
-- ========================================

-- 显示所有已安装的扩展
SELECT 
    extname AS extension_name,
    extversion AS version,
    nspname AS schema
FROM pg_extension e
JOIN pg_namespace n ON e.extnamespace = n.oid
WHERE extname NOT IN ('plpgsql')
ORDER BY extname;

-- 显示自定义函数
SELECT 
    proname AS function_name,
    pg_get_function_result(oid) AS return_type,
    pronargs AS argument_count
FROM pg_proc
WHERE pronamespace = (SELECT oid FROM pg_namespace WHERE nspname = 'public')
    AND proname IN ('get_database_size_mb', 'get_table_sizes', 'get_connection_stats', 'get_slow_queries', 'database_health_check')
ORDER BY proname;

EOSQL

echo "=========================================="
echo "RedFire PostgreSQL扩展安装完成！"
echo "已安装的扩展："
echo "  - timescaledb: 时序数据库扩展"
echo "  - uuid-ossp: UUID生成扩展"
echo "  - pg_stat_statements: 查询统计扩展"
echo "  - intarray: 数组操作扩展"
echo "  - btree_gin/btree_gist: 索引扩展"
echo "  - bloom: 布隆过滤器扩展"
echo "  - fuzzystrmatch: 模糊匹配扩展"
echo "  - tablefunc: 表函数扩展"
echo "  - pgcrypto: 加密扩展"
echo "  - postgis: 地理信息扩展 (如果可用)"
echo "  - pg_trgm: 三元组匹配扩展 (如果可用)"
echo ""
echo "已创建的自定义函数："
echo "  - get_database_size_mb(): 获取数据库大小"
echo "  - get_table_sizes(): 获取表大小统计"
echo "  - get_connection_stats(): 获取连接统计"
echo "  - get_slow_queries(): 获取慢查询统计"
echo "  - database_health_check(): 数据库健康检查"
echo "=========================================="
